package com.videoanalysis.recognizer;

import ai.djl.Application;
import ai.djl.Device;
import ai.djl.MalformedModelException;
import ai.djl.inference.Predictor;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelNotFoundException;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.TranslateException;
import com.videoanalysis.RecognitionResult;

import java.io.IOException;
import java.nio.file.Path;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class LocalImageRecognizer implements AutoCloseable, com.videoanalysis.recognizer.ImageRecognizer {
    private ZooModel<Image, Classifications> model;
    private Predictor<Image, Classifications> predictor;

    public LocalImageRecognizer() {
        try {
            // 使用预训练的ResNet-50模型
            Criteria<Image, Classifications> criteria = Criteria.builder()
                    .setTypes(Image.class, Classifications.class)
                    .optArtifactId("ai.djl.pytorch:resnet")
                    .optProgress(new ProgressBar())
                    .optOption("mapLocation", "true")
                    .build();

            model = ModelZoo.loadModel(criteria);

            predictor = model.newPredictor();
        } catch (ModelNotFoundException | MalformedModelException | IOException e) {
            throw new RuntimeException("Failed to load image recognition model", e);
        }
    }

    /**
     * 批量识别多个图像中的内容
     *
     * @param imagePaths 图像文件路径列表
     * @return 每个图像的识别结果映射，key为图像路径，value为识别结果列表
     */
    @Override
    public Map<Path, List<RecognitionResult>> recognizeImages(List<Path> imagePaths) {
        Map<Path, List<RecognitionResult>> results = new HashMap<>();

        for (Path imagePath : imagePaths) {
            try {
                Image img = ImageFactory.getInstance().fromFile(imagePath);
                Classifications classifications = predictor.predict(img);

                // 获取前5个最可能的结果并转换为RecognitionResult
                List<RecognitionResult> imageResults = classifications.topK(5)
                        .stream()
                        .map(c -> new RecognitionResult(c.getClassName(), c.getProbability()))
                        .toList();

                results.put(imagePath, imageResults);
            } catch (IOException | TranslateException e) {
                // 记录错误但继续处理其他图像
                results.put(imagePath, Collections.emptyList());
                throw new RuntimeException("Failed to recognize image: " + imagePath, e);
            }
        }

        return results;
    }

    /**
     * 关闭模型和预测器，释放资源
     */
    public void close() {
        if (predictor != null) {
            predictor.close();
        }
        if (model != null) {
            model.close();
        }
    }

    /**
     * 获取模型信息
     *
     * @return 模型描述信息
     */
    public String getModelInfo() {
        if (model != null) {
            return String.format("Model: %s\nDevice: %s",
                    model.getName(),
                    Device.cpu());
        }
        return "Model not loaded";
    }
}
